#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/1/23 20:02
# @Author  : Cxk

import os
import h5py
import numpy as np
from vgg16_keras import VGGNet
import multiprocessing

model = VGGNet()

def get_imlist(path):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')]
def child_process(img_path):
    child_norm_feat = model.vgg_extract_feat(img_path)
    child_img_name = os.path.split(img_path)[1]
    return child_norm_feat, child_img_name

if __name__ == "__main__":
    database = 'img'
    index = 'models/vgg_featureCNN.h5'
    img_list = get_imlist(database)
    feats = []
    names = []
    pool = multiprocessing.Pool(processes = 10) #多进程处理 10个
    for i, img_path in enumerate(img_list):
        try:
            norm_feat,img_name= pool.apply(child_process, (img_path, ))   #维持执行的进程总数为processes
            # 当一个进程执行完毕后会添加新的进程进去
            feats.append(norm_feat)
            names.append(img_name)
            print("extracting feature from image No. %d , %d images in total" % ((i + 1), len(img_list)))
        except:
            continue
    pool.close()
    pool.join()   #调用join之前，先调用close函数
    # 否则会出错。执行完close后不会有新的进程加入到pool,join函数等待所有子进程结束
    feats = np.array(feats)

    output = index
    h5f = h5py.File(output, 'w')
    h5f.create_dataset('dataset_1', data=feats)
    h5f.create_dataset('dataset_2', data=np.string_(names))
    h5f.close()
